#!/usr/bin/python
import sys, string
import numpy as np
import matplotlib.pyplot as plt

def get_klens(kpts, b1, b2):
    klens = []
    for i in range(len(kpts)-1):
        klens.append( np.linalg.norm(kpts[i,0]*b1 + kpts[i,1]*b2 - (kpts[i+1,0]*b1 + kpts[i+1,1]*b2)))
    return klens

def read_new(folder_name,nk):
  Eig = []
  for ik in range(nk):
    fp = open(folder_name + "ener_k" + str(ik+1) + "_s1.dat", "r")
    lines = fp.readlines()
    fp.close()
    tmp = []
    for i in range(1,len(lines)):
      w = lines[i].split()
      tmp.append(eval(w[1]))
    Eig.append(tmp)
  return Eig
def read_orig(out_file):
  fp = open(out_file,"r")
  lines = fp.readlines()
  fp.close()

  Eig = []
  for i in range(len(lines)):
    if "number of k points=" in lines[i]:
      w = lines[i].split()
      nk = eval(w[4])
    if "number of Kohn-Sham states" in lines[i]:
      w = lines[i].split()
      nbnd = eval(w[4])
      if nbnd%8 == 0:
        nlines = int(nbnd/8)
      else:
        nlines = int(nbnd/8) + 1
    if "bands (ev):" in lines[i]:
      tmp = []
      for j in range(nlines):
        w = lines[i+j+2].split()
        tmp = tmp + w
      for it in range(len(tmp)):
        tmp[it]  = eval(tmp[it])
      Eig.append(tmp)
  return Eig,nk,nbnd

nplots = 2
#Eig_pzr,nk,nbnd = read_orig("ONCVPZ_relax/out")
#Eig_pber,nk,nbnd = read_orig("ONCVPBE_relax/out")
#Eig_Mo,nk,nbnd_Mo = read_orig("out_SR")
Eig_W,nk,nbnd_W = read_orig("out")
#Eig_uspzr,nk,nbnd = read_orig("USoftPZ_relax/out")
#Eig_uspzx,nk,nbnd = read_orig("USoftPZ_3p28/out")
#Sz_kb = np.load("Sz_kb.npy")
EfMo = 3.0002
EfW = 3.0248

kpts = np.array([[0.0,0.0,0.0], [0.5,0.0,0.0], [2/3.,1/3., 0.0], [0.0,0.0,0.0]])
b1, b2 = np.array([1.000000, -0.577350,  0.000000]), np.array([0.000000,  1.154701,  0.000000])
klens = get_klens(kpts, b1, b2)
print(klens)

kpath = np.linspace(0, klens[0], 180)
kpath = np.append(kpath, np.linspace(klens[0],klens[1] + klens[0], 180))
kpath = np.append(kpath, np.linspace(klens[1] + klens[0],klens[1] + klens[2] + klens[0], 181))

#Eig_pzr = np.array(Eig_pzr) - Efpzr
#Eig_Mo = np.array(Eig_Mo) - EfMo
Eig_W = np.array(Eig_W) - EfW
#Eig_pber = np.array(Eig_pber) - Efpber
#Eig_uspzr = np.array(Eig_uspzr) - Efuspzr
#Eig_uspzx = np.array(Eig_uspzx) - Efuspzr
#Eig_AB_u = np.array(Eig_AB_u) - Ef
nplots = 1
fig, axarr = plt.subplots(1,nplots)
axarr = [axarr]
lwi = 2.2
#K = range(nk)
#for i in range(nbnd_Mo):
#  axarr[0].plot(K,Eig_Mo[:,i],color = 'C0',linewidth=lwi,linestyle = '-')
for i in range(nbnd_W):
  axarr[0].plot(kpath,Eig_W[:,i],color = 'k',linewidth=lwi,linestyle = '-')

print(len(kpath), len(Eig_W[:,0]))

#np.save("Eig_Mo", Eig_Mo)
#np.save("Eig_W", Eig_W)

yl = -1
yu = 2
fs = 22
for i in range(nplots):
  axarr[i].set_ylim(yl,yu)
  axarr[i].set_xlim(0,sum(klens))
  axarr[i].tick_params(labelsize=18,width=2.5, length = 7)
  axarr[i].set_xticks([0, klens[0], klens[0] + klens[1], sum(klens)])
  axarr[i].axvline(klens[0],yl,yu, linewidth = 2.4, color = 'k')
  axarr[i].axvline(klens[0] + klens[1],yl,yu, linewidth = 2.4, color = 'k')
  #axarr[i].plot([klens[0],],[yl,yu],linewidth=2.0,color='k')
  #axarr[i].plot([360,360],[yl,yu],linewidth=2.0,color='k')
  #axarr[i].set_xticks([0])
  axarr[i].set_xticklabels(["$\Gamma$","$\mathrm{M}$","$\mathrm{K}$","$\Gamma$"],fontsize = 18)

#axarr[2].set_title('ABp',fontsize = fs)
#axarr[3].set_title('AB',fontsize = fs)

axarr[0].set_ylabel("Energy (eV)",fontsize = 20)
#axarr[0].set_title("SR", fontsize = 18)
#axarr[1].set_title("FR", fontsize = 18)
#axarr[2].set_title(r"B$^{Se/Mo}$", fontsize = 20)

for axis in ['top','bottom','left','right']:
  for i in range(nplots):
    axarr[i].spines[axis].set_linewidth(3.0)
#plt.ylim(-0.2,1.25)

plt.tight_layout()
plt.savefig("Panel_0.png",dpi=400)
plt.show()
